#!/usr/bin/env python
import numpy as np
from maputility2 import containers
from maputility2.skyutils import fastOpen

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='Healpix Arithmatic Tools')
    parser.add_argument("f1", type=str, help='file1')
    parser.add_argument("f2", type=str, help='file2')
    parser.add_argument("fout", type=str, help='fileout = file1 ops file2')
    parser.add_argument("ops", type=str, help='operators +|-|*|/|M')
    parser.add_argument("-zero", type=float, default=1.e-50, help='absolute value < zero are set to zero, only in ops div')
    parser.add_argument("-nowcs", action='store_true', default=False,
            help='do not check the wcs infos')
    args = parser.parse_args()

    if args.nowcs:
        m1 = fastOpen(args.f1, verbose=False, container=containers.PixMap)
        m2 = fastOpen(args.f2, verbose=False, container=containers.PixMap)
    else:
        m1 = fastOpen(args.f1, verbose=False)
        m2 = fastOpen(args.f2, verbose=False)

    if args.ops == '+':
        m1 += m2
        m1.writeHistory('{} + {}'.format(args.f1, args.f2), origin='HpxArith')
    elif args.ops == '-':
        m1 -= m2
        m1.writeHistory('{} - {}'.format(args.f1, args.f2), origin='HpxArith')
    elif args.ops == '*':
        m1 *= m2
        m1.writeHistory('{} * {}'.format(args.f1, args.f2), origin='HpxArith')
    elif args.ops == '/':
        m1 /= m2
        m1.writeHistory('{} / {}'.format(args.f1, args.f2), origin='HpxArith')
        zero_fil = (~np.isfinite(m1.mapdata)) | (np.abs(m1.mapdata)<args.zero)
        m1.mapdata[zero_fil] = 0.
    elif args.ops.lower() == 'm':
        if m2.nlayers == 1:
            m1._mask_exec(m2.map, allmask=True)
        else:
            m1._mask_exec(m2.mapdata)
        m1.writeHistory('masked with {}'.format(args.f2), origin='HpxArith')
    else:
        raise NotImplementedError('Unsupported operator ({}), please use "+|-|*|/|M"')
    m1.saveMap(args.fout)
